{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2018 Google LLC  \n",
    "  \n",
    " Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    " you may not use this file except in compliance with the License.  \n",
    " You may obtain a copy of the License at  \n",
    "  \n",
    "     http://www.apache.org/licenses/LICENSE-2.0  \n",
    "  \n",
    " Unless required by applicable law or agreed to in writing, software  \n",
    " distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    " See the License for the specific language governing permissions and  \n",
    " limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the dependency .py files, if any.\n",
    "! git clone https://github.com/GoogleCloudPlatform/cloudml-samples.git\n",
    "! cp cloudml-samples/tpu/templates/tpu_gan_estimator/* .\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import numpy as np\n",
    "import os\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = 5\n",
    "OUTPUT_DIM = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator_fn(generator_inputs):\n",
    "    outputs = tf.layers.dense(generator_inputs, OUTPUT_DIM)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator_fn(data, generator_inputs):\n",
    "    outputs = tf.layers.dense(data, 1)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_model_fn(features, labels, mode, params):\n",
    "    # build model\n",
    "    global_step = tf.train.get_global_step()\n",
    "\n",
    "    generator_inputs = features\n",
    "    real_data = labels\n",
    "\n",
    "    with tf.variable_scope('shared', reuse=tf.AUTO_REUSE):\n",
    "        gan_model = tf.contrib.gan.gan_model(generator_fn, discriminator_fn, real_data, generator_inputs)\n",
    "\n",
    "    predictions = gan_model.generated_data\n",
    "    loss = None\n",
    "    train_op = None\n",
    "\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        # define loss\n",
    "        gan_loss = tf.contrib.gan.gan_loss(gan_model, add_summaries=False)\n",
    "        loss = gan_loss.generator_loss\n",
    "\n",
    "        # define train_op\n",
    "        optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "        dummy_optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "\n",
    "        # wrapper to make the optimizer work with TPUs\n",
    "        if params['use_tpu']:\n",
    "            optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
    "\n",
    "        gan_train_ops = tf.contrib.gan.gan_train_ops(gan_model, gan_loss, optimizer, dummy_optimizer)\n",
    "\n",
    "        # tf.contrib.gan's train op does not manage global steps in it\n",
    "        train_op = tf.group(gan_train_ops.generator_train_op, global_step.assign_add(1))\n",
    "\n",
    "    if params['use_tpu']:\n",
    "        # TPU version of EstimatorSpec\n",
    "        return tf.contrib.tpu.TPUEstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op)\n",
    "    else:\n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dis_model_fn(features, labels, mode, params):\n",
    "    # build model\n",
    "    global_step = tf.train.get_global_step()\n",
    "\n",
    "    generator_inputs = features\n",
    "    real_data = labels\n",
    "\n",
    "    with tf.variable_scope('shared', reuse=tf.AUTO_REUSE):\n",
    "        gan_model = tf.contrib.gan.gan_model(generator_fn, discriminator_fn, real_data, generator_inputs)\n",
    "\n",
    "    predictions = {\n",
    "        'discriminator_gen_outputs': gan_model.discriminator_gen_outputs,\n",
    "        'discriminator_real_outputs': gan_model.discriminator_real_outputs}\n",
    "    loss = None\n",
    "    train_op = None\n",
    "\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        # define loss\n",
    "        gan_loss = tf.contrib.gan.gan_loss(gan_model, add_summaries=False)\n",
    "        loss = gan_loss.discriminator_loss\n",
    "\n",
    "        # define train_op\n",
    "        optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "        dummy_optimizer = tf.train.RMSPropOptimizer(learning_rate=0.05)\n",
    "\n",
    "        # wrapper to make the optimizer work with TPUs\n",
    "        if params['use_tpu']:\n",
    "            optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
    "\n",
    "        gan_train_ops = tf.contrib.gan.gan_train_ops(gan_model, gan_loss, dummy_optimizer, optimizer)\n",
    "\n",
    "        # tf.contrib.gan's train op does not manage global steps in it\n",
    "        train_op = tf.group(gan_train_ops.discriminator_train_op, global_step.assign_add(1))\n",
    "\n",
    "    if params['use_tpu']:\n",
    "        # TPU version of EstimatorSpec\n",
    "        return tf.contrib.tpu.TPUEstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op)\n",
    "    else:\n",
    "        return tf.estimator.EstimatorSpec(\n",
    "            mode=mode,\n",
    "            predictions=predictions,\n",
    "            loss=loss,\n",
    "            train_op=train_op)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_input_fn(params={}):\n",
    "    # make some fake noise\n",
    "    data_size = 100\n",
    "    noise_tensor = tf.random_normal((data_size, INPUT_DIM))\n",
    "    real_data_tensor = tf.random_uniform((data_size, OUTPUT_DIM))\n",
    "\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((noise_tensor, real_data_tensor))\n",
    "    dataset = dataset.repeat().shuffle(10)\n",
    "\n",
    "    # TPUEstimator passes params when calling input_fn\n",
    "    batch_size = params.get('train_batch_size', 16)\n",
    "    dataset = dataset.batch(batch_size, drop_remainder=True)\n",
    "\n",
    "    # TPUs need to know all dimensions when the graph is built\n",
    "    # Datasets know the batch size only when the graph is run\n",
    "    def set_shapes(features, labels):\n",
    "        features_shape = features.get_shape().merge_with([batch_size, None])\n",
    "        labels_shape = labels.get_shape().merge_with([batch_size, None])\n",
    "\n",
    "        features.set_shape(features_shape)\n",
    "        labels.set_shape(labels_shape)\n",
    "\n",
    "        return features, labels\n",
    "\n",
    "    dataset = dataset.map(set_shapes)\n",
    "    dataset = dataset.prefetch(tf.contrib.data.AUTOTUNE)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "    # pass the args as params so the model_fn can use\n",
    "    # the TPU specific args\n",
    "    params = vars(args)\n",
    "\n",
    "    if args.use_tpu:\n",
    "        # additional configs required for using TPUs\n",
    "        tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(args.tpu)\n",
    "        tpu_config = tf.contrib.tpu.TPUConfig(\n",
    "            num_shards=8, # using Cloud TPU v2-8\n",
    "            iterations_per_loop=args.save_checkpoints_steps)\n",
    "\n",
    "        # use the TPU version of RunConfig\n",
    "        gen_config = tf.contrib.tpu.RunConfig(\n",
    "            cluster=tpu_cluster_resolver,\n",
    "            model_dir=os.path.join(args.model_dir, 'generator'),\n",
    "            tpu_config=tpu_config,\n",
    "            save_checkpoints_steps=args.save_checkpoints_steps,\n",
    "            save_summary_steps=100)\n",
    "\n",
    "        dis_config = tf.contrib.tpu.RunConfig(\n",
    "            cluster=tpu_cluster_resolver,\n",
    "            model_dir=os.path.join(args.model_dir, 'discriminator'),\n",
    "            tpu_config=tpu_config,\n",
    "            save_checkpoints_steps=args.save_checkpoints_steps,\n",
    "            save_summary_steps=100)\n",
    "\n",
    "        # TPUEstimator\n",
    "        gen_estimator = tf.contrib.tpu.TPUEstimator(\n",
    "            model_fn=gen_model_fn,\n",
    "            config=gen_config,\n",
    "            params=params,\n",
    "            train_batch_size=args.train_batch_size,\n",
    "            eval_batch_size=32,\n",
    "            export_to_tpu=False)\n",
    "\n",
    "        dis_estimator = tf.contrib.tpu.TPUEstimator(\n",
    "            model_fn=dis_model_fn,\n",
    "            config=dis_config,\n",
    "            params=params,\n",
    "            train_batch_size=args.train_batch_size,\n",
    "            eval_batch_size=32,\n",
    "            export_to_tpu=False)\n",
    "    else:\n",
    "        gen_config = tf.estimator.RunConfig(model_dir=os.path.join(args.model_dir, 'generator'))\n",
    "        dis_config = tf.estimator.RunConfig(model_dir=os.path.join(args.model_dir, 'discriminator'))\n",
    "\n",
    "        gen_estimator = tf.estimator.Estimator(\n",
    "            gen_model_fn,\n",
    "            config=gen_config,\n",
    "            params=params)\n",
    "\n",
    "        dis_estimator = tf.estimator.Estimator(\n",
    "            dis_model_fn,\n",
    "            config=dis_config,\n",
    "            params=params)\n",
    "\n",
    "    # manage the training loop\n",
    "    for _ in range(3):\n",
    "        tf.logging.info('Training Discriminator')\n",
    "        dis_estimator.train(train_input_fn, steps=100)\n",
    "\n",
    "        tf.logging.info('Training Generator')\n",
    "        gen_estimator.train(train_input_fn, steps=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "\n",
    "parser.add_argument(\n",
    "    '--model-dir',\n",
    "    type=str,\n",
    "    default='/tmp/tpu-template',\n",
    "    help='Location to write checkpoints and summaries to.  Must be a GCS URI when using Cloud TPU.')\n",
    "parser.add_argument(\n",
    "    '--train-batch-size',\n",
    "    type=int,\n",
    "    default=16,\n",
    "    help='The training batch size.  The training batch is divided evenly across the TPU cores.')\n",
    "parser.add_argument(\n",
    "    '--save-checkpoints-steps',\n",
    "    type=int,\n",
    "    default=100,\n",
    "    help='The number of training steps before saving each checkpoint.')\n",
    "parser.add_argument(\n",
    "    '--use-tpu',\n",
    "    action='store_true',\n",
    "    help='Whether to use TPU.')\n",
    "parser.add_argument(\n",
    "    '--tpu',\n",
    "    default=None,\n",
    "    help='The name or GRPC URL of the TPU node.  Leave it as `None` when training on CMLE.')\n",
    "\n",
    "args, _ = parser.parse_known_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO(user): change this\n",
    "args.model_dir = 'gs://your-gcs-bucket'\n",
    "\n",
    "# Get hostname from environment using ipython magic.\n",
    "# This returns a list.\n",
    "hostname = !hostname\n",
    "\n",
    "args.tpu = hostname[0]\n",
    "args.use_tpu = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use gcloud command line tool to create a TPU in the same zone as the VM instance.\n",
    "! gcloud compute tpus create `hostname` \\\n",
    "  --zone `gcloud compute instances list --filter=\"name=$(hostname)\" --format 'csv[no-heading](zone)'`\\\n",
    "  --network default \\\n",
    "  --range 10.101.1.0 \\\n",
    "  --version 1.13\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use gcloud command line tool to delete the TPU.\n",
    "! gcloud compute tpus delete `hostname` \\\n",
    "  --zone `gcloud compute instances list --filter=\"name=$(hostname)\" --format 'csv[no-heading](zone)'`\\\n",
    "  --quiet\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
